# -*- coding: utf-8 -*-
"""
@Time    : 2025/2/27 19:42 
@Author  : ZhangShenao 
@File    : apply.py 
@Desc    : 应用与评估模型
"""
from typing import List

from gensim.models import Word2Vec


def similarity_search(query: str, top_n: int = 5) -> List[str]:
    """
    基于语义相似度搜索
    :param query: 搜索的文本
    :param top_n: 返回的相似度最高的前n个结果
    :return: 返回的相似度最高的前top_n个结果
    """

    # 加载模型
    model = Word2Vec.load("./word2vec.model")
    print("模型加载完成")

    # 使用模型
    # 生成query对应的词向量
    print(f"query: {query}")
    print(f"词向量: \n{model.wv[query]}")

    # 找到最相似的词
    similar_words = model.wv.most_similar(query, topn=top_n)
    return similar_words
